# -*- coding: utf-8 -*-
import or1200_definitions
import sys
sys.path.append("..")
import pyrtl
from  pyrtl import GPUSim
from or1200_mem2reg import Or1200Mem2Reg
from or1200_reg2mem import Or1200Reg2Mem

dw = or1200_definitions.OR1200_OPERAND_WIDTH
aw = or1200_definitions.OR1200_REGFILE_ADDR_WIDTH

class Or1200Lsu(object):
    def __init__(self):
        self.addrbase = pyrtl.Input(bitwidth=32, name='lsu_addrbase')
        self.addrofs = pyrtl.Input(bitwidth=32, name='lsu_addrofs')
        self.lsu_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH, name='lsu_lsu_op')
        self.lsu_datain = pyrtl.Input(bitwidth=dw, name='lsu_lsu_datain')
        self.lsu_dataout = pyrtl.Output(bitwidth=dw, name='lsu_lsu_dataout')
        # self.lsu_dataout = pyrtl.WireVector(bitwidth=dw)
        self.lsu_stall = pyrtl.Output(bitwidth=1,name='lsu_stall')
        self.lsu_unstall_out = pyrtl.Output(bitwidth=1,name='lsu_lsu_unstall_out')
        self.lsu_unstall = pyrtl.WireVector(bitwidth=1)
        self.du_stall = pyrtl.Input(bitwidth=1, name='lsu_du_stall')
        self.except_align = pyrtl.WireVector(bitwidth=1)
        self.except_align_out = pyrtl.Output(bitwidth=1,name='lsu_except_align_out')

        self.except_dbuserr = pyrtl.Output(bitwidth=1,name='lsu_except_dbuserr')


        self.dcpu_adr_o = pyrtl.WireVector(bitwidth=32)
        self.dcpu_cycstb_o = pyrtl.WireVector(bitwidth=1)
        self.dcpu_adr_o_out = pyrtl.Output(bitwidth=32,name='lsu_dcpu_adr_o_out')
        self.dcpu_cycstb_o_out = pyrtl.Output(bitwidth=1,name='lsu_dcpu_cycstb_o_out')
        self.dcpu_we_o = pyrtl.Output(bitwidth=1,name='lsu_dcpu_we_o')
        #self.dcpu_we_o = pyrtl.WireVector(bitwidth=1)
        # self.dcpu_sel_o = pyrtl.Output(bitwidth=4)
        self.dcpu_tag_o = pyrtl.Output(bitwidth=4,name='lsu_dcpu_tag_o')
        #self.dcpu_tag_o = pyrtl.WireVector(bitwidth=4)
        self.dcpu_dat_o = pyrtl.Output(bitwidth=32,name='lsu_dcpu_dat_o')
        #self.dcpu_dat_o = pyrtl.WireVector(bitwidth=32)
        self.dcpu_dat_i = pyrtl.Input(bitwidth=32, name='lsu_dcpu_dat_i')
        self.dcpu_ack_i = pyrtl.Input(bitwidth=1, name='lsu_dcpu_ack_i')
        self.dcpu_rty_i = pyrtl.Input(bitwidth=1, name='lsu_dcpu_rty_i')
        self.dcpu_err_i = pyrtl.Input(bitwidth=1, name='lsu_dcpu_err_i')
        self.dcpu_tag_i = pyrtl.Input(bitwidth=4, name='lsu_dcpu_tag_i')
        
        self.dcpu_sel_o = pyrtl.Output(bitwidth=4,name='lsu_dcpu_sel_o')


        self.lsu_unstall_out<<=self.lsu_unstall
        self.except_align_out<<=self.except_align
        self.dcpu_cycstb_o_out<<= self.dcpu_cycstb_o
        self.dcpu_adr_o_out<<=self.dcpu_adr_o


        
        self.lsu_stall <<= self.dcpu_rty_i & self.dcpu_cycstb_o
        self.lsu_unstall <<= self.dcpu_ack_i
        self.except_align <<= ((self.lsu_op == or1200_definitions.OR1200_LSUOP_SH) | (self.lsu_op == or1200_definitions.OR1200_LSUOP_LHZ) | (self.lsu_op == or1200_definitions.OR1200_LSUOP_LHS)) & self.dcpu_adr_o[0]\
        | ((self.lsu_op == or1200_definitions.OR1200_LSUOP_SW) | (self.lsu_op == or1200_definitions.OR1200_LSUOP_LWZ) | (self.lsu_op == or1200_definitions.OR1200_LSUOP_LWS)) & or1200_definitions.Or(self.dcpu_adr_o[0:2])
        self.except_dbuserr <<= self.dcpu_err_i & (self.dcpu_tag_i == or1200_definitions.OR1200_DTAG_BE)
        
        self.dcpu_adr_o <<= self.addrbase + self.addrofs
        self.dcpu_cycstb_o <<= pyrtl.select(self.du_stall | self.lsu_unstall | self.except_align, pyrtl.Const(0b0, bitwidth=1), or1200_definitions.Or(self.lsu_op))
        self.dcpu_we_o <<= self.lsu_op[3]
        self.dcpu_tag_o <<=  pyrtl.select(self.dcpu_cycstb_o, or1200_definitions.OR1200_DTAG_ND, or1200_definitions.OR1200_DTAG_IDLE)
        
        #initialize all the parts
        general_lsu = GeneralLsu()
        
        general_lsu.lsu_op <<= self.lsu_op
        general_lsu.dcpu_adr_o <<= self.dcpu_adr_o
        self.dcpu_sel_o <<= general_lsu.dcpu_sel_o
        
        mem2reg = Or1200Mem2Reg()
        #addr作为一个模块中的输入信号，无法被赋值
        mem2reg.addr <<= self.dcpu_adr_o[0:2]
        mem2reg.lsu_op <<= self.lsu_op
        mem2reg.memdata <<= self.dcpu_dat_i
        self.lsu_dataout <<= mem2reg.regdata
        
        reg2mem = Or1200Reg2Mem()
        reg2mem.addr <<= self.dcpu_adr_o[0:2]
        reg2mem.lsu_op <<= self.lsu_op
        reg2mem.regdata <<= self.lsu_datain
        self.dcpu_dat_o <<= reg2mem.memdata


class Or1200Lsu_wire(object):
    def __init__(self):
        self.addrbase = pyrtl.WireVector(bitwidth=32, name='lsu_lsu_addrbase')
        self.addrofs = pyrtl.WireVector(bitwidth=32, name='lsu_lsu_addrofs')
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH, name='lsu_lsu_op')
        self.lsu_datain = pyrtl.WireVector(bitwidth=dw, name='lsu_lsu_datain')
        self.lsu_dataout = pyrtl.WireVector(bitwidth=dw, name='lsu_lsu_dataout')
        # self.lsu_dataout = pyrtl.WireVector(bitwidth=dw)
        self.lsu_stall = pyrtl.WireVector(bitwidth=1, name='lsu_lsu_stall')
        self.lsu_unstall_out = pyrtl.WireVector(bitwidth=1, name='lsu_lsu_unstall_out')
        self.lsu_unstall = pyrtl.WireVector(bitwidth=1)
        self.du_stall = pyrtl.WireVector(bitwidth=1, name='lsu_du_stall')
        self.except_align = pyrtl.WireVector(bitwidth=1)
        self.except_align_out = pyrtl.WireVector(bitwidth=1, name='lsu_except_align_out')

        self.except_dbuserr = pyrtl.WireVector(bitwidth=1, name='lsu_except_dbuserr')

        self.dcpu_adr_o = pyrtl.WireVector(bitwidth=32)
        self.dcpu_cycstb_o = pyrtl.WireVector(bitwidth=1)
        self.dcpu_adr_o_out = pyrtl.WireVector(bitwidth=32, name='lsu_dcpu_adr_o_out')
        self.dcpu_cycstb_o_out = pyrtl.WireVector(bitwidth=1, name='lsu_dcpu_cycstb_o_out')
        self.dcpu_we_o = pyrtl.WireVector(bitwidth=1, name='lsu_dcpu_we_o')
        # self.dcpu_we_o = pyrtl.WireVector(bitwidth=1)
        # self.dcpu_sel_o = pyrtl.WireVector(bitwidth=4)
        self.dcpu_tag_o = pyrtl.WireVector(bitwidth=4, name='lsu_dcpu_tag_o')
        # self.dcpu_tag_o = pyrtl.WireVector(bitwidth=4)
        self.dcpu_dat_o = pyrtl.WireVector(bitwidth=32, name='lsu_dcpu_dat_o')
        # self.dcpu_dat_o = pyrtl.WireVector(bitwidth=32)
        self.dcpu_dat_i = pyrtl.WireVector(bitwidth=32, name='lsu_dcpu_dat_i')
        self.dcpu_ack_i = pyrtl.WireVector(bitwidth=1, name='lsu_dcpu_ack_i')
        self.dcpu_rty_i = pyrtl.WireVector(bitwidth=1, name='lsu_dcpu_rty_i')
        self.dcpu_err_i = pyrtl.WireVector(bitwidth=1, name='lsu_dcpu_err_i')
        self.dcpu_tag_i = pyrtl.WireVector(bitwidth=4, name='lsu_dcpu_tag_i')

        self.dcpu_sel_o = pyrtl.WireVector(bitwidth=4, name='lsu_dcpu_sel_o')

        self.lsu_unstall_out <<= self.lsu_unstall
        self.except_align_out <<= self.except_align
        self.dcpu_cycstb_o_out <<= self.dcpu_cycstb_o
        self.dcpu_adr_o_out <<= self.dcpu_adr_o

        self.lsu_stall <<= self.dcpu_rty_i & self.dcpu_cycstb_o
        self.lsu_unstall <<= self.dcpu_ack_i
        self.except_align <<= ((self.lsu_op == or1200_definitions.OR1200_LSUOP_SH) | (
                    self.lsu_op == or1200_definitions.OR1200_LSUOP_LHZ) | (
                                           self.lsu_op == or1200_definitions.OR1200_LSUOP_LHS)) & self.dcpu_adr_o[0] \
                              | ((self.lsu_op == or1200_definitions.OR1200_LSUOP_SW) | (
                    self.lsu_op == or1200_definitions.OR1200_LSUOP_LWZ) | (
                                             self.lsu_op == or1200_definitions.OR1200_LSUOP_LWS)) & or1200_definitions.Or(
            self.dcpu_adr_o[0:2])
        self.except_dbuserr <<= self.dcpu_err_i & (self.dcpu_tag_i == or1200_definitions.OR1200_DTAG_BE)

        self.dcpu_adr_o <<= self.addrbase + self.addrofs
        self.dcpu_cycstb_o <<= pyrtl.select(self.du_stall | self.lsu_unstall | self.except_align,
                                            pyrtl.Const(0b0, bitwidth=1), or1200_definitions.Or(self.lsu_op))
        self.dcpu_we_o <<= self.lsu_op[3]
        self.dcpu_tag_o <<= pyrtl.select(self.dcpu_cycstb_o, or1200_definitions.OR1200_DTAG_ND,
                                         or1200_definitions.OR1200_DTAG_IDLE)

        # initialize all the parts
        general_lsu = GeneralLsu()

        general_lsu.lsu_op <<= self.lsu_op
        general_lsu.dcpu_adr_o <<= self.dcpu_adr_o
        self.dcpu_sel_o <<= general_lsu.dcpu_sel_o

        mem2reg = Or1200Mem2Reg()
        # addr作为一个模块中的输入信号，无法被赋值
        mem2reg.addr <<= self.dcpu_adr_o[0:2]
        mem2reg.lsu_op <<= self.lsu_op
        mem2reg.memdata <<= self.dcpu_dat_i
        self.lsu_dataout <<= mem2reg.regdata

        reg2mem = Or1200Reg2Mem()
        reg2mem.addr <<= self.dcpu_adr_o[0:2]
        reg2mem.lsu_op <<= self.lsu_op
        reg2mem.regdata <<= self.lsu_datain
        self.dcpu_dat_o <<= reg2mem.memdata


class GeneralLsu(object):
    def __init__(self):
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.dcpu_adr_o = pyrtl.WireVector(bitwidth=32)
        self.dcpu_sel_o = pyrtl.WireVector(bitwidth=4)
        
        with pyrtl.conditional_assignment:
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1000, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBZ, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1000, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBS, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1000, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b01, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBZ, pyrtl.Const(0b01, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBS, pyrtl.Const(0b01, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0010, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBZ, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0010, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBS, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0010, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b11, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0001, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBZ, pyrtl.Const(0b11, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0001, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LBS, pyrtl.Const(0b11, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0001, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SH, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LHZ, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LHS, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1100, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SH, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0011, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LHZ, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0011, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LHS, pyrtl.Const(0b10, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b0011, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_SW, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1111, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LWZ, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1111, bitwidth=4)
            with pyrtl.concat(self.lsu_op, self.dcpu_adr_o[0:2]) == pyrtl.concat(or1200_definitions.OR1200_LSUOP_LWS, pyrtl.Const(0b00, bitwidth=2)):
                self.dcpu_sel_o |= pyrtl.Const(0b1111, bitwidth=4)

if __name__ == '__main__':
    lsu = Or1200Lsu()
    input_width = []

    #com = pyrtl.CompiledSimulation()
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])
    import random

    max_ = 0
    com_time = 0
    ess_time = 0
    num = 1
    for i in range(num):
        print(i)

        inputs_random = {}
        inputs_select = {}
        inputs_old = {}
        inputs_ = {}

        for inp in input_width:
            inputs_random[inp[0]] = []
            inputs_select[inp[0]] = []
            result = random.randint(0, 2 ** inp[1] - 1)
            inputs_old[inp[0]] = result

        for j in range(131072):
            for inp in input_width:
                result = random.randint(0, 2 ** inp[1] - 1)
                inputs_[inp[0]] = result
                inputs_random[inp[0]].append(result)
            if random.random() < 0.2:
                for name in inputs_old:
                    inputs_select[name].append(inputs_old[name])
            else:
                for name in inputs_:
                    inputs_select[name].append(inputs_[name])
                    inputs_old[name] = inputs_[name]
        #time1 = com.step_multiple(inputs_select)
    sim = pyrtl.GPUSim_now.GPUSim('lsu')
    sim.make_input(inputs_select)
